package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/streadway/amqp"
)

type Client struct {
	conn          *websocket.Conn
	lastHeartbeat time.Time
}

var (
	wsClients       = make(map[*Client]bool)
	wsMutex         sync.Mutex
	rabbitConn      *amqp.Connection
	rabbitCh        *amqp.Channel
	rabbitConnected bool
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func wsHandler(c *gin.Context) {
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Println("upgrade error:", err)
		return
	}
	client := &Client{
		conn:          conn,
		lastHeartbeat: time.Now(),
	}
	wsMutex.Lock()
	wsClients[client] = true
	wsMutex.Unlock()
	log.Println("ws client connected:", conn.RemoteAddr())

	go loopMsg(client)
}

func loopMsg(client *Client) {
	defer func() {
		wsMutex.Lock()
		delete(wsClients, client)
		wsMutex.Unlock()
		client.conn.Close()
		log.Println("ws client disconnected:", client.conn.RemoteAddr())
	}()
	for {
		_, msg, err := client.conn.ReadMessage()
		if err != nil {
			log.Println("Read message error:", err)
			return
		}
		var data map[string]string
		if err := json.Unmarshal(msg, &data); err != nil {
			log.Println("Receive non-json data:", string(msg))
			continue
		}
		switch data["type"] {
		case "ping":
			client.lastHeartbeat = time.Now()
			pongMsg, _ := json.Marshal(map[string]string{"type": "pong"})
			client.conn.WriteMessage(websocket.TextMessage, pongMsg)
		case "connect":
			ip := data["ip"]
			port := data["port"]
			username := data["username"]
			password := data["password"]
			url := fmt.Sprintf("amqp://%s:%s@%s:%s/", username, password, ip, port)
			conn, err := amqp.Dial(url)
			if err != nil {
				sendToClient(client, map[string]string{"type": "error", "message": "Connect to RabbitMQ failed: " + err.Error()})
				return
			}
			ch, err := conn.Channel()
			if err != nil {
				sendToClient(client, map[string]string{"type": "error", "message": "Create RabbitMQ channel failed: " + err.Error()})
				return
			}
			rabbitConn = conn
			rabbitCh = ch
			rabbitConnected = true
			sendToClient(client, map[string]string{"type": "info", "message": "Connected to RabbitMQ"})
		case "disconnect":
			if rabbitConnected {
				rabbitCh.Close()
				rabbitConn.Close()
				rabbitConnected = false
				sendToClient(client, map[string]string{"type": "info", "message": "Disconnected from RabbitMQ"})
			}
		case "send":
			if !rabbitConnected {
				sendToClient(client, map[string]string{"type": "error", "message": "Not connected to RabbitMQ"})
				continue
			}
			topic := data["topic"]
			payload := data["payload"]
			err := rabbitCh.Publish("", topic, false, false, amqp.Publishing{
				ContentType: "text/plain",
				Timestamp:   time.Now(),
				Body:        []byte(payload),
			})
			if err != nil {
				sendToClient(client, map[string]string{"type": "error", "message": "Send message failed: " + err.Error()})
				continue
			}
			sendToClient(client, map[string]string{"type": "info", "message": "Message sent to topic " + topic})
		}
	}
}

func sendToClient(client *Client, data map[string]string) {
	msgBytes, _ := json.Marshal(data)
	client.conn.WriteMessage(websocket.TextMessage, msgBytes)
}

func main() {
	r := gin.Default()
	r.Static("/", "./wwwroot")
	r.GET("/ws", wsHandler)
	r.Run(":8080")
}
